from typing import Any, Dict

import torch


class Summarizer:
    """
    总结器是评估器(Validator)的核心成员，用于记录评估过程中的状态，并基于这些状态计算生成评估指标。
    评估指标能够反映模型精度。有两方面的评估指标，分别为损失(loss)和其它自定义指标(metric)。
    若无自定义指标，则可仅采用损失作为评估标准，以确定模型的拟合状态——欠拟合/恰拟合/过拟合。
    """

    def __init__(self):
        self._init_metric_cache()
        self.loss_dict, self.n_batches = None, 0

    def update(self, loss_dict, results, labels):
        # 更新评估指标缓存和损失缓存
        self._update_metrics(results, labels)
        if self.loss_dict is None:
            self.loss_dict = loss_dict
        else:
            for key, value in loss_dict.items():
                if not torch.isnan(value):
                    self.loss_dict[key] += value
        self.n_batches += 1

    def summarize(self):
        avg_loss_dict = dict()
        for k, v in self.loss_dict.items():
            avg_loss_dict[k] = v / self.n_batches
        self.loss_dict, self.n_batches = None, 0
        return avg_loss_dict, self._summarize_metrics()

    def _update_metrics(self, results, labels):
        # 更新缓存，被更新的缓存可用于最终计算评估指标
        for result, label in zip(results, labels):
            self._update_metric_cache(result, label)

    def _summarize_metrics(self):
        metrics = self._calculate_metrics()
        self._init_metric_cache()
        return metrics

    def _init_metric_cache(self):
        """
        metric缓存的内容因具体任务而定，清空的操作不像loss缓存那么固定。
        因此开放接口，以供子类依情况进行清空操作

        Example:
            >>> self.tf_form = []  # 分类评估时，要初始化一个true/false列表
        """
        pass

    def _update_metric_cache(self, result, label):
        """
        更新metric缓存

        Example:
            >>> # 评估分类模型时，需要将预测结果与gt做对比，生成true/false flag，并缓存
            >>> self.tf_form.append(result.item() == label[0])
        """
        pass

    def _calculate_metrics(self) -> Dict[str, Any]:
        """
        计算自定义评估指标，该指标不同于评估损失，
        不仅能用于判断拟合状态，还可用于做同类任务下的不同算法实现的横向对比

        Example:
            >>> # 评估分类模型时，只需统计t/f flag表中t的比例即可得到acc
            >>> import numpy as np
            >>> return dict(acc=np.array(self.tf_form).mean())

        """
        pass
